package cn.mengfly.testnrz.curator;

import org.apache.zookeeper.Watcher;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 基于Zookeeper的分布式锁的实现
 *
 * @author Mengfly
 */
public class ZkLock implements Lock {
    private static final String ZK_PATH = "/test/lock";
    private static final String LOCK_PREFIX = ZK_PATH + "/";
    private static final long WAIT_TIME = 3;

    private String lockedShortPath = null;
    private String lockedPath = null;
    private String priorPath = null;

    private final AtomicInteger lockCount = new AtomicInteger();
    private Thread thread;

    public ZkLock() {
        if (!ZkClient.inst().isNodeExists(ZK_PATH)) {
            ZkClient.inst().createNode(ZK_PATH, null);
        }
    }

    @Override
    public boolean lock() {
        // 可重入，确保同一线程可以重复加锁
        synchronized (this) {
            if (lockCount.get() == 0) {
                thread = Thread.currentThread();
                lockCount.incrementAndGet();
            } else {
                if (!thread.equals(Thread.currentThread())) {
                    return false;
                }
                lockCount.incrementAndGet();
                return true;
            }
            try {
                boolean locked = tryLock();
                if (locked) {
                    return true;
                }
                while (!locked) {
                    // 等待
                    await();
                    // 获取等待的子节点列表
                    List<String> waiters = getWaiters();

                    // 判断是否需加锁成功
                    if (checkLocked(waiters)) {
                        locked = true;
                    }
                }
                return true;
            } catch (Exception e) {
                e.printStackTrace();
                unlock();
            }
            return false;
        }
    }

    /**
     * 尝试加锁
     *
     * @return 是否加锁成功
     */
    private boolean tryLock() throws Exception {
        lockedPath = ZkClient.inst().createEphemeralSeqNode(LOCK_PREFIX);
        if (null == lockedPath) {
            throw new Exception("zk error");
        }
        lockedShortPath = getShortPath(lockedPath);

        // 获取加锁的队列
        List<String> waiters = getWaiters();

        // 获取等待的子节点列表，判断自己是否是第一个
        if (checkLocked(waiters)) {
            return true;
        }

        // 判断自己排在第几个
        int index = Collections.binarySearch(waiters, lockedShortPath);
        if (index < 0) {
            // 网络抖动，获取到的子节点列表里可能已经没有自己了
            throw new Exception("节点没有找到：" + lockedShortPath);
        }
        // 如果自己没有获取到锁
        priorPath = LOCK_PREFIX + waiters.get(index - 1);
        return false;
    }

    private boolean checkLocked(List<String> waiters) {
        // 节点按照编号，升序排序
        Collections.sort(waiters);

        // 如果是第一个，代表自己已经获得了锁
        if (lockedShortPath.equals(waiters.get(0))) {
            System.out.println("成功获取到分布式锁， 节点为 " + lockedShortPath);
            return true;
        }
        return false;
    }

    private List<String> getWaiters() {
        return ZkClient.inst().getChildren(ZK_PATH);
    }

    /**
     * 等待
     * 监听前一个节点的删除事件
     *
     * @throws Exception 可能会有Zk异常，网络异常
     */
    private void await() throws Exception {
        if (null == priorPath) {
            throw new Exception("prior path error");
        }
        final CountDownLatch latch = new CountDownLatch(1);
        // Watcher 一次性订阅
        Watcher w = watchedEvent -> {
            latch.countDown();
        };
        // 开始监听
        ZkClient.inst().getClient().getData().usingWatcher(w).forPath(priorPath);
        // 限时等待，最长加锁时间为3s
        latch.await(WAIT_TIME, TimeUnit.SECONDS);
    }

    private String getShortPath(String lockedPath) {
        int index = lockedPath.lastIndexOf(LOCK_PREFIX);
        if (index >= 0) {
            index += ZK_PATH.length() + 1;
            return index <= lockedPath.length() ? lockedPath.substring(index) : "";
        }
        return null;
    }

    @Override
    public boolean unlock() {
        // 只有加锁的线程能够解锁
        if (!thread.equals(Thread.currentThread())) {
            return false;
        }
        // 减少可重入的计数
        int newLockCount = lockCount.decrementAndGet();
        if (newLockCount < 0) {
            throw new IllegalMonitorStateException("计数不对：" + lockedPath);
        }
        // 如果计数不为0，直接返回
        if (newLockCount != 0) {
            return true;
        }
        // 删除临时节点
        return ZkClient.inst().delete(lockedPath);
    }
}
